Passare da Modalità immediata di PyTorch a Triton richiede uno spostamento dal vedere i tensori come oggetti monolitici a vederli come insiemi di blocchi discreti e gestibili blocchi o matrici.
1. Tensori PyTorch vs. Triton
È fondamentale distinguere tensori Triton da tensori PyTorch. Un tensore PyTorch è un oggetto Python lato host che contiene forma, tipo, dispositivo, stride e metadati di archiviazione. Al contrario, Triton opera con i puntatori grezzi ai dati all'interno di blocchi di memoria specifici, consentendo ottimizzazioni a livello molto più basso.
2. Il collo di bottiglia dell'esecuzione immediata
Nell'esecuzione standard immediata, ogni operazione (ad esempio, addizione poi ReLU) richiede un lancio separato del kernel e un viaggio completo nella memoria globale. Questo rappresenta il principale collo di bottiglia nel calcolo moderno con GPU. Triton supera questo problema attraverso l'unione di operazioni all'interno di un singolo kernel che elabora blocchi di dati (ad esempio, 128, 256 o 512 elementi) direttamente nella memoria interna al chip.
3. Il paradigma basato sui blocchi
Invece del pensiero a livello scalare dei thread CUDA, Triton utilizza SPMD (Programma singolo, Dati multipli) a livello di blocco. Scrivi un solo kernel, e Triton avvia molteplici istanze su una griglia. Ogni istanza utilizza il suo program_id per calcolare quale "blocco" di memoria possiede.
4. Configurazione dell'ambiente
Per iniziare, installa Triton in un ambiente pulito (usando Conda o venv) per assicurarti che non ci siano conflitti di dipendenze con gli strumenti CUDA esistenti: pip install triton.